package com.tuniu.agents.common.logger;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.agent.SelectionAgent;
import com.tuniu.agents.message.MessageLogger;
import com.tuniu.agents.message.MessageSender;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;

@Aspect
@Component
public class SelectionAgentLoggerAspect extends LoggerAspect{

    private static final Logger LOGGER = LoggerFactory.getLogger(SelectionAgentLoggerAspect.class);

    @Autowired
    private ChatEnvironmentMemory chatEnvironmentMemory;

    public SelectionAgentLoggerAspect(ApplicationContext applicationContext, MessageSender messageSender) {
        super(applicationContext, messageSender);
    }

    // Define pointcut: all call and stream methods of classes that inherit from Agent
    @Pointcut("execution(* com.tuniu.agents.*.call(String, String)) || " +
            "execution(* com.tuniu.agents..*.call(String, java.util.List)) || " +
            "execution(* com.tuniu.agents..*.call(String, java.util.List, java.util.Map)) || " +
            "execution(* com.tuniu.agents..*.callResponse(String, String)) || " +
            "execution(* com.tuniu.agents..*.callResponse(String, java.util.List)) || " +
            "execution(* com.tuniu.agents..*.callResponse(String, java.util.List, java.util.Map)) || " +
            "execution(* com.tuniu.agents..*.stream(String, String)) || " +
            "execution(* com.tuniu.agents..*.stream(String, java.util.List)) || " +
            "execution(* com.tuniu.agents..*.stream(String, java.util.List, java.util.Map)) || " +
            "execution(* com.tuniu.agents..*.streamResponse(String, String)) || " +
            "execution(* com.tuniu.agents..*.streamResponse(String, java.util.List)) || " +
            "execution(* com.tuniu.agents..*.streamResponse(String, java.util.List, java.util.Map))"
    )
    public void agentMethods() {}

    /**
     * Define pointcut: only intercept call and stream methods of classes that inherit from SelectionAgent/AnalysisAgent
     */
    @Pointcut("agentMethods() && (within(com.tuniu.agents.common.agent.SelectionAgent+))")
    public void agentClasses() {}

    @Around("agentClasses()")
    public Object aroundAgentMethod(ProceedingJoinPoint joinPoint) throws Throwable {

        // Get the target object (instance of SelectionAgent) and store the last selected ID in a global variable
        SelectionAgent target = (SelectionAgent) joinPoint.getTarget();
        String conversationId = this.getConversationId(joinPoint.getArgs());
        if(StringUtils.isNotBlank(conversationId)) {
            String selectionAgentId = target.getAgentId();
            final Collection<String> main_intent_history = (Collection<String>) Optional.ofNullable(this.chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.MAIN_INTENT_HISTORY)).orElse(new ArrayList<>()); // Historical main intent
            main_intent_history.add(selectionAgentId);
            this.chatEnvironmentMemory.add(conversationId, Map.of(
                    AgentsConstants.AgentsConversationEnvironmentKey.LAST_SELECTION_AGENT_ID, selectionAgentId,
                    AgentsConstants.AgentsConversationEnvironmentKey.MAIN_INTENT_HISTORY, main_intent_history));
        }

        // Now get the value of a specific field (e.g., 'messageLogger' field)
        String fieldName = "messageLogger"; // Replace with the field name you need to access
        MessageLogger logger = null;
        try {
            Field field = target.getClass().getSuperclass().getDeclaredField(fieldName);
            field.setAccessible(true); // Set field to be accessible
            logger = (MessageLogger) field.get(target); // Get field value
        } catch (NoSuchFieldException | IllegalAccessException e) {
            // Log error or throw custom exception
            LOGGER.error("Failed to access field: " + fieldName);
        }
        Object result = getResult(joinPoint, logger);

        return result;
    }

    @Override
    public String getConversationId(Object[] args) {
        Object conversationId = args[0];
        return conversationId == null ? null : String.valueOf(conversationId);
    }

}
